#!/usr/bin/env python3
import os
import sys
import time
import datetime
import signal


# 颜色输出
class Colors:
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    ENDC = '\033[0m'


# 默认配置
DEFAULT_PID = None  # None 表示监控脚本自身
DEFAULT_INTERVAL = 2  # 刷新间隔(秒)
DEFAULT_COUNT = None  # 刷新次数(None表示无限循环)


def get_process_info(pid):
    """获取进程的基本信息"""
    try:
        with open(f"/proc/{pid}/stat", 'r') as f:
            stat = f.read().split()

        with open(f"/proc/{pid}/status", 'r') as f:
            status = f.read()

        # 解析CPU时间
        utime = int(stat[13])  # 用户态CPU时间
        stime = int(stat[14])  # 内核态CPU时间
        start_time = int(stat[21])  # 进程启动时间

        # 解析内存使用
        vmsize = 0
        rss = 0
        for line in status.split('\n'):
            if line.startswith('VmSize:'):
                vmsize = int(line.split()[1])
            elif line.startswith('VmRSS:'):
                rss = int(line.split()[1])

        # 解析进程名称
        comm = stat[1].strip('()')

        # 获取系统时钟滴答数和CPU核心数
        clock_ticks = os.sysconf(os.sysconf_names['SC_CLK_TCK'])
        cpu_count = os.cpu_count()

        return {
            'pid': pid,
            'name': comm,
            'utime': utime,
            'stime': stime,
            'start_time': start_time,
            'vmsize': vmsize,  # KB
            'rss': rss,  # KB
            'clock_ticks': clock_ticks,
            'cpu_count': cpu_count
        }
    except Exception as e:
        print(f"无法获取进程 {pid} 的信息: {e}")
        return None


def get_threads_info(pid):
    """获取进程所有线程的CPU使用情况"""
    threads = []
    try:
        # 获取所有线程ID
        if os.path.exists(f"/proc/{pid}/task"):
            tids = [int(tid) for tid in os.listdir(f"/proc/{pid}/task") if tid.isdigit()]

            for tid in tids:
                try:
                    with open(f"/proc/{pid}/task/{tid}/stat", 'r') as f:
                        stat = f.read().split()

                    # 解析线程CPU时间
                    utime = int(stat[13])  # 用户态CPU时间
                    stime = int(stat[14])  # 内核态CPU时间

                    # 解析线程名称
                    comm = stat[1].strip('()')

                    threads.append({
                        'tid': tid,
                        'name': comm,
                        'utime': utime,
                        'stime': stime
                    })
                except:
                    continue
    except:
        pass

    return threads


def calculate_cpu_usage(old_info, new_info, interval):
    """计算CPU使用率"""
    if not old_info or not new_info:
        return 0.0

    # 计算总CPU时间差
    total_time = ((new_info['utime'] - old_info['utime']) +
                  (new_info['stime'] - old_info['stime']))

    # 转换为秒
    total_time_sec = total_time / new_info['clock_ticks']

    # 计算CPU使用率百分比
    cpu_percent = (total_time_sec / interval) * 100 * new_info['cpu_count']

    return min(cpu_percent, 100.0 * new_info['cpu_count'])


def calculate_thread_cpu_usage(old_threads, new_threads, interval, clock_ticks, cpu_count):
    """计算线程CPU使用率"""
    thread_usage = []

    # 创建TID到线程信息的映射
    old_thread_map = {t['tid']: t for t in old_threads}

    for new_thread in new_threads:
        tid = new_thread['tid']
        if tid in old_thread_map:
            old_thread = old_thread_map[tid]

            # 计算总CPU时间差
            total_time = ((new_thread['utime'] - old_thread['utime']) +
                          (new_thread['stime'] - old_thread['stime']))

            # 转换为秒
            total_time_sec = total_time / clock_ticks

            # 计算CPU使用率百分比
            cpu_percent = (total_time_sec / interval) * 100 * cpu_count

            thread_usage.append({
                'tid': tid,
                'name': new_thread['name'],
                'cpu_percent': min(cpu_percent, 100.0)
            })

    # 按CPU使用率排序
    thread_usage.sort(key=lambda x: x['cpu_percent'], reverse=True)
    return thread_usage


def print_process_info(info, cpu_percent):
    """打印进程信息"""
    pid = info['pid']
    name = info['name']
    vmsize = info['vmsize'] / 1024  # MB
    rss = info['rss'] / 1024  # MB

    # 根据CPU使用率设置颜色
    if cpu_percent > 80:
        cpu_color = Colors.RED
    elif cpu_percent > 50:
        cpu_color = Colors.YELLOW
    else:
        cpu_color = Colors.GREEN

    print(f"PID: {pid}")
    print(f"进程名称: {name}")
    print(f"CPU使用率: {cpu_color}{cpu_percent:.2f}%{Colors.ENDC}")
    print(f"内存使用: RSS={rss:.2f} MB, VmSize={vmsize:.2f} MB")
    print("-" * 50)


def print_threads_info(threads):
    """打印线程信息"""
    print("线程信息:")
    print(f"{'TID':<8} {'名称':<15} {'CPU%'}")
    print("-" * 30)

    for thread in threads[:20]:  # 只显示前20个线程
        tid = thread['tid']
        name = thread['name'][:15]
        cpu_percent = thread['cpu_percent']

        # 根据CPU使用率设置颜色
        if cpu_percent > 50:
            cpu_color = Colors.RED
        elif cpu_percent > 20:
            cpu_color = Colors.YELLOW
        else:
            cpu_color = Colors.GREEN

        print(f"{tid:<8} {name:<15} {cpu_color}{cpu_percent:.2f}%{Colors.ENDC}")


def main():
    """主函数"""
    # 使用默认配置
    pid = DEFAULT_PID
    interval = DEFAULT_INTERVAL
    count = DEFAULT_COUNT

    # 如果未指定PID，默认监控脚本自身
    if pid is None:
        pid = os.getpid()
        print(f"未指定PID，默认监控当前脚本 (PID: {pid})")
        time.sleep(1)

    # 检查进程是否存在
    if not os.path.exists(f"/proc/{pid}"):
        print(f"错误: 进程 {pid} 不存在")
        sys.exit(1)

    # 首次采样
    old_process_info = get_process_info(pid)
    old_threads_info = get_threads_info(pid)

    # 设置退出信号处理
    def signal_handler(sig, frame):
        print("\n监控已停止")
        sys.exit(0)

    signal.signal(signal.SIGINT, signal_handler)

    # 主循环
    iteration = 0
    try:
        while count is None or iteration < count:
            # 清屏
            os.system('clear')
            # 再次采样
            new_process_info = get_process_info(pid)
            new_threads_info = get_threads_info(pid)

            if not new_process_info:
                print(f"进程 {pid} 已终止")
                break

            # 计算CPU使用率
            cpu_percent = calculate_cpu_usage(old_process_info, new_process_info, interval)

            # 计算线程CPU使用率
            thread_usage = calculate_thread_cpu_usage(
                old_threads_info, new_threads_info,
                interval, new_process_info['clock_ticks'],
                new_process_info['cpu_count']
            )

            # 打印信息
            print(f"监控时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
            print(f"监控进程: {pid}")
            print("=" * 50)

            print_process_info(new_process_info, cpu_percent)
            print_threads_info(thread_usage)

            # 底部提示
            print("\n按 Ctrl+C 停止监控")
            print(f"刷新间隔: {interval}秒")
            # 等待指定的间隔
            time.sleep(interval)
            # 更新旧信息
            old_process_info = new_process_info
            old_threads_info = new_threads_info

            iteration += 1

    except KeyboardInterrupt:
        print("\n监控已停止")


if __name__ == "__main__":
    main()